package top.jolyoulu.sql

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql._
import org.apache.spark.{SparkConf, SparkContext}

/**
 * @Author: JolyouLu
 * @Date: 2024/5/16 19:21
 * @Description
 */
object Spark01_SparkSQL_UDAF_Old_Generics {
  def main(args: Array[String]): Unit = {
    //创建SparkSQL的运行环境
    val sparkSQL: SparkConf = new SparkConf().setMaster("local[*]").setAppName("SparkSQL")
    val spark = SparkSession.builder().config(sparkSQL).getOrCreate()
    val sc: SparkContext = spark.sparkContext
    import spark.implicits._ //将spark转换规则引入
    //DataFrame转换
    val path: String = this.getClass.getClassLoader.getResource("datas/user.json").toURI.getPath
    val df: DataFrame = spark.read.json(path)
    //转换位DataSet
    val ds: Dataset[User] = df.as[User]
    //Spark 3.x之前的聚合函数使用
    val column: TypedColumn[User, Long] = new MyAvgUDAF().toColumn
    ds.select(column).show()

    //关闭环境
    spark.close()
  }
  //ds的数据结构
  case class User(username:String,age:Long)
  //自定义聚合函数类，计算年龄的平均值
  //IN：输入的数据类型
  //BUF：中间存储数据类型
  //OUT：输出的数据类型
  case class Buff(var total:Long,var count:Long)
  class MyAvgUDAF extends Aggregator[User,Buff,Long]{
    //初始值
    override def zero: Buff = {
      Buff(0L,0L)
    }
    //根据输入的数据更新缓冲区的数据
    override def reduce(buff: Buff, in: User): Buff = {
      buff.total = buff.total + in.age
      buff.count = buff.count + 1
      buff
    }
    //合并缓冲区
    override def merge(buff1: Buff, buff2: Buff): Buff = {
      buff1.total = buff1.total + buff2.total
      buff1.count = buff1.count + buff2.count
      buff1
    }
    //计算结果
    override def finish(reduction: Buff): Long = {
      reduction.total / reduction.count
    }
    //缓冲区数据编码逻辑
    override def bufferEncoder: Encoder[Buff] = Encoders.product
    //缓冲区数据解码逻辑
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }
}
